from quantsim import QHeap
from math import pi, sqrt

class Qmos(QHeap):
    
    def __init__(self, size):
        super(Qmos, self).__init__(size)
        self.size = self._QHeap__size
        self.old_solutions = set()
    
    def initialize(self):
        self.reset()
        size = self.size
        for index in range(size):
            self.mix(index)
            
    def apply_oracle(self, oracle):
        # oracle returns True if index is a solution and False otherwise
        
        for index in range(2**self.size):
            if oracle(index):
                self._QHeap__states[index] *= -1 # Flip the amplitude
                
    def apply_d(self):
        # calculate mean
        states = self._QHeap__states
        total = 0
        for index in range(2**self.size):
            total += states[index]
        mean = total / float(2**self.size)
        
        for index in range(2**self.size):
            distance_from_mean = mean - states[index]
            states[index] = mean + distance_from_mean
         
    def count_solutions(self, oracle):
        # Not sure how to implement a true quantum counting
        # algorithm. This is a stand-in
        
        total_solutions = 0
        for index in range(2**self.size):
            if oracle(index):
                total_solutions += 1
        return total_solutions
    
    def measure_all_qubits(self):
        total = 0
        for index in range(self.size):
            total += self.measure(index) * 2 ** index
        return total
    
    def quantum_search(self, oracle):
        num_solutions = self.count_solutions(oracle)
        def augmented_oracle(index):
            return oracle(index) and index not in self.old_solutions
        if num_solutions == 0:
            return set()
        num_iterations = int(pi/4 * sqrt(2**self.size / float(num_solutions)))
        for i in range(num_solutions):
            self.initialize()
            for j in range(num_iterations):
                self.apply_oracle(augmented_oracle)
                self.apply_d()
            solution = self.measure_all_qubits()
            self.old_solutions.add(solution)
        return self.old_solutions
    
if __name__ == "__main__":
    
    def oracle(index):
        return index == 0 or index == 2 or index == 30
    
    qmos = Qmos(10)
    print qmos.quantum_search(oracle)
    